"""
pip install pycrypto

使用 PyCrypto 进行 AES/ECB/PKCS#5(7) 加密
http://likang.me/blog/2013/06/05/python-pycrypto-aes-ecb-pkcs-5/
"""

import base64
import json

from Crypto.Cipher import AES


def pkcs7_unpad(s):
    """ 拿掉最后补位的字符
    加密数据的最后一个字节用来存补位字符数量
    """
    return s[:-ord(s[len(s)-1:])]


def pkcs7_pad(s: str):
    s = s.encode('utf-8')
    pad_len = AES.block_size - len(s) % AES.block_size
    # s += bytes([pad_len]) * pad_len
    s += chr(pad_len).encode('utf-8') * pad_len
    return s


def encrypt(data, aes_key: str, iv: str, mode: int = AES.MODE_CBC):
    cipher = AES.new(aes_key, mode, iv)
    data_str = pkcs7_pad(json.dumps(data))
    encrypted = cipher.encrypt(data_str)
    return encrypted


def decrypt(encrypted, aes_key: str, iv: str, mode: int = AES.MODE_CBC):
    """
    Args:
        encrypted: str, 加密数据
        aes_key: str, 会话密钥（16 字节）
        iv: str, 初始向量（16 字节）
        mode: int, 解密算法 AES-128-CBC，PKCS#7 填充
    """
    cipher = AES.new(aes_key, mode, iv)
    decrypted = json.loads(pkcs7_unpad(cipher.decrypt(encrypted)))
    return decrypted


def main():
    aes_key = 'key:'
    iv = 'iv:'
    data = {'nihao': 123}

    aes_key += '0' * (24 - len(aes_key))  # 16, 24, 32
    iv += '0' * (16 - len(iv))  # 16
    encrypted = encrypt(data, aes_key, iv)
    print(encrypted)
    print(decrypt(encrypted, aes_key, iv))


if __name__ == '__main__':
    main()
